#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
Author: Sjy
Description: 自定义二维图显示组件
'''

__all__ = ["CustomPlot"]
import sys
from PyQt5 import QtWidgets
from PyQt5.QtWidgets import QGroupBox, QVBoxLayout, QApplication
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
import matplotlib as mpl
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar

# 设置中文字体
def set_chinese_font():
    """设置matplotlib中文字体"""
    try:
        # 优先使用微软雅黑
        plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
    except:
        try:
            # 如果没有微软雅黑，尝试使用其他中文字体
            plt.rcParams['font.sans-serif'] = ['SimHei', 'SimSun', 'KaiTi', 'FangSong', 'STSong']
        except:
            print("警告：未找到合适的中文字体，可能无法正确显示中文")
    
    # 解决负号显示问题
    plt.rcParams['axes.unicode_minus'] = False

class CustomPlot(QGroupBox):
    """
    自定义二维图显示组件
    
    该组件继承自QGroupBox，提供了一个带有标题的二维图显示区域。
    使用matplotlib进行绘图，支持实时数据更新和样式自定义。
    
    Attributes:
        figure (matplotlib.figure.Figure): matplotlib图形对象
        canvas (FigureCanvas): matplotlib画布对象
        ax (matplotlib.axes.Axes): 坐标轴对象
    """
    
    def __init__(self, title="", parent=None):
        """
        初始化自定义二维图组件
        
        Args:
            title (str): GroupBox的标题
            parent (QWidget, optional): 父组件. Defaults to None.
        """
        super(CustomPlot, self).__init__(title, parent)
        # 设置中文字体
        set_chinese_font()
        self._init_ui()
        
    def _init_ui(self):
        """初始化UI组件"""
        # 创建布局
        layout = QVBoxLayout()
        layout.setContentsMargins(5, 5, 5, 5)
        
        # 创建matplotlib图形
        self.figure = plt.figure(figsize=(8, 6))
        self.canvas = FigureCanvas(self.figure)
        self.ax = self.figure.add_subplot(111)
        
        # 设置默认样式
        self.ax.grid(True)
        self.ax.set_xlabel('X轴')
        self.ax.set_ylabel('Y轴')
        
        # 添加画布到布局
        layout.addWidget(self.canvas)
        self.setLayout(layout)
        
        # 启用鼠标事件
        self.canvas.mpl_connect('scroll_event', self._on_scroll)
        self.canvas.mpl_connect('button_press_event', self._on_press)
        self.canvas.mpl_connect('button_release_event', self._on_release)
        self.canvas.mpl_connect('motion_notify_event', self._on_motion)
        
        # 初始化平移变量
        self._panning = False
        self._last_x = None
        self._last_y = None
        
    def _on_scroll(self, event):
        """处理鼠标滚轮事件"""
        if event.inaxes:
            # 获取当前视图范围
            cur_xlim = self.ax.get_xlim()
            cur_ylim = self.ax.get_ylim()
            
            # 计算缩放中心
            xdata = event.xdata
            ydata = event.ydata
            
            # 计算缩放比例
            base_scale = 1.1
            scale = base_scale if event.button == 'up' else 1/base_scale
            
            # 计算新的范围
            new_width = (cur_xlim[1] - cur_xlim[0]) / scale
            new_height = (cur_ylim[1] - cur_ylim[0]) / scale
            
            # 设置新的范围
            relx = (cur_xlim[1] - xdata)/(cur_xlim[1] - cur_xlim[0])
            rely = (cur_ylim[1] - ydata)/(cur_ylim[1] - cur_ylim[0])
            
            self.ax.set_xlim([xdata - new_width*(1-relx), xdata + new_width*relx])
            self.ax.set_ylim([ydata - new_height*(1-rely), ydata + new_height*rely])
            
            self.canvas.draw()
            
    def _on_press(self, event):
        """处理鼠标按下事件"""
        if event.button == 1:  # 左键
            self._panning = True
            self._last_x = event.xdata
            self._last_y = event.ydata
            
    def _on_release(self, event):
        """处理鼠标释放事件"""
        if event.button == 1:  # 左键
            self._panning = False
            
    def _on_motion(self, event):
        """处理鼠标移动事件"""
        if self._panning and event.inaxes:
            dx = event.xdata - self._last_x
            dy = event.ydata - self._last_y
            
            # 更新视图范围
            cur_xlim = self.ax.get_xlim()
            cur_ylim = self.ax.get_ylim()
            
            self.ax.set_xlim(cur_xlim[0] - dx, cur_xlim[1] - dx)
            self.ax.set_ylim(cur_ylim[0] - dy, cur_ylim[1] - dy)
            
            self._last_x = event.xdata
            self._last_y = event.ydata
            
            self.canvas.draw()
        
    def plot(self, x_data, y_data, label=None, color=None, line_style='-'):
        """
        绘制数据
        
        Args:
            x_data (array-like): X轴数据
            y_data (array-like): Y轴数据
            label (str, optional): 图例标签. Defaults to None.
            color (str, optional): 线条颜色. Defaults to None.
            line_style (str, optional): 线条样式. Defaults to '-'.
        """
        self.ax.plot(x_data, y_data, label=label, color=color, linestyle=line_style)
        self.ax.legend()
        self.canvas.draw()
        
    def scatter(self, x_data, y_data, label=None, color=None, marker='o'):
        """
        绘制散点图
        
        Args:
            x_data (array-like): X轴数据
            y_data (array-like): Y轴数据
            label (str, optional): 图例标签. Defaults to None.
            color (str, optional): 点颜色. Defaults to None.
            marker (str, optional): 点样式. Defaults to 'o'.
        """
        self.ax.scatter(x_data, y_data, label=label, color=color, marker=marker)
        self.ax.legend()
        self.canvas.draw()
        
    def clear(self):
        """清除当前图形"""
        self.ax.clear()
        self.ax.grid(True)
        self.canvas.draw()
        
    def set_title(self, title):
        """
        设置图形标题
        
        Args:
            title (str): 图形标题
        """
        self.ax.set_title(title)
        self.canvas.draw()
        
    def set_labels(self, x_label, y_label):
        """
        设置坐标轴标签
        
        Args:
            x_label (str): X轴标签
            y_label (str): Y轴标签
        """
        self.ax.set_xlabel(x_label)
        self.ax.set_ylabel(y_label)
        self.canvas.draw()
        
    def set_limits(self, x_min=None, x_max=None, y_min=None, y_max=None):
        """
        设置坐标轴范围
        
        Args:
            x_min (float, optional): X轴最小值. Defaults to None.
            x_max (float, optional): X轴最大值. Defaults to None.
            y_min (float, optional): Y轴最小值. Defaults to None.
            y_max (float, optional): Y轴最大值. Defaults to None.
        """
        if x_min is not None and x_max is not None:
            self.ax.set_xlim(x_min, x_max)
        if y_min is not None and y_max is not None:
            self.ax.set_ylim(y_min, y_max)
        self.canvas.draw()
        
    def set_theme(self, theme='default'):
        """
        设置图形主题
        
        Args:
            theme (str, optional): 主题名称. Defaults to 'default'.
        """
        plt.style.use(theme)
        self.canvas.draw() 


if __name__ == "__main__":
    app = QApplication(sys.argv)
    main_window = QtWidgets.QDialog()
    main_window.setWindowTitle("数据可视化")
    main_window.resize(800, 600)

    import numpy as np

    # 创建组件
    plot_widget = CustomPlot(title="数据可视化")
    

    layout = QVBoxLayout()
    main_window.setLayout(layout)
    layout.addWidget(plot_widget)

    # 生成示例数据
    x = np.linspace(0, 10, 100)
    y = np.sin(x)

    # 绘制数据
    plot_widget.plot(x, y, label="正弦曲线", color="blue")
    # 设置标题和标签
    plot_widget.set_title("正弦函数图像")
    plot_widget.set_labels("时间", "振幅")

    # 设置坐标轴范围
    plot_widget.set_limits(x_min=0, x_max=10, y_min=-1.5, y_max=1.5)


    main_window.show()
    sys.exit(app.exec_())

